﻿using System;
using ASquared;

namespace iwpm.optimization
{
    public class LagrangeMultiplierMatrix : Matrix
    {
        private static double s_defaultVal = 0.0; 

        private double _alpha = 1;
        private Matrix _prevLagr = new Matrix();
        private Matrix _prevLagr2 = new Matrix();
        private int _prevIteration = int.MaxValue;
        private double _tolerance = 0.001;
        private int _maxIterations = 100;
        private int _innerIteration = 0;
        private bool _useAugmented = false;

        public double Tolerance
        {
            get
            {
                return _tolerance;
            }
            set
            {
                _tolerance = value;
            }
        }

        public int Iteration
        {
            get
            {
                return _innerIteration;
            }
        }

        public int MaxIterations
        {
            get
            {
                return _maxIterations;
            }
            set
            {
                _maxIterations = value;
            }
        }

        public double Alpha
        {
            get
            {
                return _alpha;
            }
            set
            {
                _alpha = value;
            }
        }

        public bool UseAugmented
        {
            get
            {
                return _useAugmented;
            }
            set
            {
                _useAugmented = value;
            }
        }

        public LagrangeMultiplierMatrix(int NumberOfMultipliers)
            : base(NumberOfMultipliers, 1, s_defaultVal)
        {
        }

        /// <summary>Updates the Lagrange Multipliers and returns whether or not it has improved</summary>
        /// <param name="iteration">The iteration of the master problem.</param>
        /// <param name="constraintValue">The values of each constraint (in matrix form).</param>
        /// <returns>Returns whether this Lagrange function has converged.</returns>
        public bool Update(int iteration, Matrix constraintValue)
        {
            if (this == null || this.table == null || this.NumOfRows == 0)
                return true;  

            // update alpha
            if (!_useAugmented)
                _alpha = 1 / (1 + 0.1 * _innerIteration);

            // update this Lagrange multiplier matrix 
            if (_prevIteration > iteration)
            {
                _innerIteration = 0;
                this.table = (new Matrix(this.NumOfRows, 1, s_defaultVal)).table;
            }
            else
            {
                _innerIteration++;
                this.table = Lagr(_alpha, constraintValue).table;
            }

            // determine convergence
            bool convg = false;
            if (_prevLagr.Length == this.Length && _prevLagr2.Length == this.Length)
            {
                convg = (_innerIteration > _maxIterations)              // hit max iterations
                    || (constraintValue.Abs() < _tolerance)             // constraint equation is below a tolerance
                    || ((_prevLagr2 - this).Magnitude() / _prevLagr.Magnitude() < _tolerance);  // the Lagrange multipliers have converged
                if (!convg)
                    convg = (_prevLagr - _prevLagr2).Abs() < _tolerance && (_prevLagr - this).Abs() < _tolerance;
            }

            // update previous information
            _prevIteration = iteration;
            _prevLagr2.table = _prevLagr.table;
            _prevLagr.table = this.table;
            
            return convg;
        }

        private Matrix Lagr(double alpha, Matrix constraintValue)
        {
            //if (_useAugmented)
            //    return ((Matrix)this + alpha * constraintValue);
            //else 
            if (constraintValue.Abs() < _tolerance)
                return this; 
            return ((Matrix)this + alpha * constraintValue / constraintValue.Magnitude());
        }

    }
}
